package com.putprize.percy

import scala.math.{log,abs,exp}

import scala.collection.mutable.{HashMap => HMap}

import org.apache.log4j.Logger

object PyotrEstimater {
  
  val _log = Logger.getLogger(this.getClass.getName)
  
  val MAX_ITERATION = 50
  val MIN_CONVERGENCE = 1e-4
  
  def run_em(ms:Map[String,Document], zs:Map[String,List[Int]],
             N:Int, M:Int,K:Int, 
             maxIter:Int=MAX_ITERATION,
             minConv:Double=MIN_CONVERGENCE):Model = {
    
    var model = PyotrModel.initModel(K,M)
    var iter = 0
    var convergence = 1.0
    var likelihood = -10000.0
    
    while (iter < maxIter && convergence > minConv) {
      iter += 1
      
      val countT = new Array[Double](K)
      val countV = new Array[Array[Double]](M)
      (0 until M).foreach(v => {
        countV(v) = new Array[Double](K)
      })
      
      var S = 0.0
      
      
      var count = 0
      
      // Inference
      _log.info("Inference ...")
      val rs = ms.par.map( x => {
        val name = x._1
        val m = x._2
        count += 1
        if (count % 1000 == 0){
          _log.info(count)
        }
        
        val res = PyotrInfer.run_vt(model,m,zs(name))
        val res1 = zs(name).map(z => (z, res._1(z))).toMap[Int,Double]
        val res2 = new Array[Map[Int,Double]](ms(name).n)
        (0 until ms(name).n).foreach(i => {
          res2(i) = zs(name).map(z => (z,res._2(i)(z))).toMap[Int,Double]
        })
        
        val res3 = res._3
        (name,(res1,res2,res3))
      })
      _log.info("Done Inference")
      
      // Estimate
      var t = 0.0 // Likelihood
      rs.foreach( r => {
        val name = r._1
        val res = r._2
        // Likelihood
        t += res._3
        // Count
        val doc = ms(name)
        (0 until doc.n).foreach( i => {
          val v = doc.vs(i)
          val c = doc.cs(i)
          zs(name).foreach( z => {
            countT(z) += c*res._2(i)(z)
            countV(v)(z) += c*res._2(i)(z)
          })
        })
        // Gamma
        var s = 0.0
        zs(name).foreach (z => {
          s += res._1(z)
          S += Gamma.iGamma(res._1(z))
        })
        S -= K*Gamma.iGamma(s) 
      })
      
      // Check
      _log.info("Likelihood "+t+" "+likelihood)
      convergence = abs( (likelihood-t)/likelihood )
      likelihood = t
      _log.info(iter + " " + likelihood + " " + convergence)
      
      // M
      model = PyotrModel.initModel(countT,countV,S,N)
      _log.info("INIT_T_VALUE "+model.INIT_T_VALUE)
    }
    model
  }
}
